# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt
import random
import time

from data_process import getUbAndLb


class PSO():
    # PSO参数设置
    def __init__(self, pN, dim, max_iter, lb=None, ub=None, particle_output=False):
        self.w = 0.8
        self.c1 = 2
        self.c2 = 2
        self.pN = pN  # 粒子数量
        self.dim = dim  # 搜索维度
        self.max_iter = max_iter  # 迭代次数
        self.X = np.random.rand(self.pN, self.dim)  # 所有粒子的位置和速度
        self.V = np.zeros_like(self.X)
        self.pbest = np.zeros((self.pN, self.dim))  # 个体经历的最佳位置和全局最佳位置
        self.gbest = np.zeros((1, self.dim))
        self.p_fit = np.zeros(self.pN)  # 每个个体的历史最佳适应值
        self.fit = 1e10  # 全局最佳适应值
        self.particle_output = particle_output
        self.binding = False
        self.EL = []  # 粒子群总体能量
        self.funcName = None
        if (lb is not None) and (ub is not None):
            assert len(lb) == len(ub), 'Lower- and upper-bounds must be the same length'
            self.lb = np.array(lb)
            self.ub = np.array(ub)
            assert np.all(ub > lb), 'All upper-bound values must be greater than lower-bound values'
            self.vhigh = np.abs(self.ub - self.lb)
            self.vlow = -self.vhigh
            # Initialize the particle's position
            self.X = self.lb + self.X * (self.ub - self.lb)
            # Initialize the particle's velocity
            self.V = self.vlow + np.random.rand(pN, dim) * (self.vhigh - self.vlow)
            self.binding = True

    # TODO 目标函数
    def getCurrentFitness(self, location):
        self.funcName = "LJ"
        return self.getLJ(location)
        # return self.getSutton_Chen(location)

    def getLJ(self, location):
        lj = 0
        for p in self.X:
            if (p == location).all():
                continue
            r = np.linalg.norm(p - location)
            lj += 4 * (r ** (-12) - r ** (-6))
        return lj

    def getSutton_Chen(self, location):
        c = 2
        a = 1
        m = 1
        sutton = 0
        p_i = 0
        p_1 = 0
        for particle in self.X:
            if (particle == location).all():
                continue
            r = np.linalg.norm(particle - location)
            p_i += (a / r) ** m
            p_1 += (a / r)
        sutton += 4 * (0.5 * p_1 - c * np.sqrt(p_i))
        return sutton

    # 初始化种群最佳位置和拟合度
    def init_Population(self):
        for i in range(self.pN):  # 因为要随机生成pN个数据，所以需要循环pN次
            self.pbest[i] = self.X[i]  # 其实就是给self.pbest定值
            tmp = self.getCurrentFitness(self.X[i])  # 得到现在最优
            self.p_fit[i] = tmp  # 这个个体历史最佳的位置
            if tmp < self.fit:  # 得到现在最优和历史最优比较大小，如果现在最优大于历史最优，则更新历史最优
                self.fit = tmp
                self.gbest = self.X[i]

    # 更新粒子位置
    def iterator(self):
        fitness = []
        for t in range(self.max_iter):  # 迭代次数，不是越多越好
            for i in range(self.pN):  # 更新gbest\pbest
                temp = self.getCurrentFitness(self.X[i])
                if temp < self.p_fit[i]:  # 更新个体最优
                    self.p_fit[i] = temp
                    self.pbest[i] = self.X[i]
                    if self.p_fit[i] < self.fit:  # 更新全局最优
                        self.gbest = self.X[i]
                        self.fit = self.p_fit[i]
            r1 = np.random.uniform(size=(self.pN, self.dim))
            r2 = np.random.uniform(size=(self.pN, self.dim))
            self.V = self.w * self.V + self.c1 * r1 * (self.pbest - self.X) + \
                     self.c2 * r2 * (self.gbest - self.X)
            self.X = self.X + self.V
            if self.binding:
                # Correct for bound violations
                maskl = self.X < self.lb
                masku = self.X > self.ub
                self.X = self.X * (~np.logical_or(maskl, masku)) + self.lb * maskl + self.ub * masku
            self.EL.append(np.sum(self.p_fit))
            fitness.append(self.fit)
            print("------------当前位置----------")
            print(self.X[0], end=" ")
            print(end="\n")
            print("------------当前最优值---------")
            print(self.fit)
        if self.particle_output:
            return fitness, self.gbest, self.p_fit, self.X
        return fitness, self.gbest


def log(fitness, bestPara, p_fit, X):
    print("*" * 50)
    print("每个粒子的当前最佳位置：")
    print('         {}'.format(X))
    print("每个粒子的当前最低势能：")
    print('         {}'.format(p_fit))
    print("粒子最佳位置：")
    print('  {}'.format(bestPara))
    print("粒子最低能量：")
    print('  {}'.format(fitness[-1]))
    print("粒子群最低能量：")
    print('  {}'.format(my_pso.EL))


def plot(fitness, my_pso, title="Au20"):
    # 画图
    plt.figure(figsize=(12, 8))
    plt.title(title + " lowest Particle energy change")
    plt.xlabel("iterators", size=14)
    plt.ylabel("fitness", size=14)
    t = np.array([t for t in range(0, ITERATION)])
    fitness = np.array(fitness)
    plt.plot(t, fitness, color='b', linewidth=3)
    plt.savefig(
        "../image/" + title + "_particle_energy" + my_pso.funcName + time.strftime("%Y-%m-%d-%H_%M_%S",
                                                                                   time.localtime()))
    plt.show()
    plt.figure(figsize=(12, 8))
    plt.title(title + " Particle swarm energy change")
    plt.xlabel("iterators", size=14)
    plt.ylabel("energy", size=14)
    plt.plot(t, my_pso.EL)
    # plt.gca().set_ylim(top=50)
    plt.ylim(-500, 500)
    plt.savefig(
        "../image/" + title + "_energy" + my_pso.funcName + time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime()))
    plt.show()


if __name__ == '__main__':
    # 程序
    ITERATION = 800
    ub, lb = getUbAndLb(path="B45-_OPT_3751")
    my_pso = PSO(pN=40, dim=3, max_iter=ITERATION, lb=lb, ub=ub, particle_output=True)
    my_pso.init_Population()
    fitness, bestPara, p_fit, X = my_pso.iterator()
    log(fitness, bestPara, p_fit, X)
    plot(fitness, my_pso, title="B_40-")
